import os.path

import yaml

from BaseParams import BaseParams
from constants import *
from data_process.DatasetGenerator import DatasetGenerator
from implement.model_distillation import LipReading


class SingleDopplerParams(BaseParams):
    def __init__(self):
        super(SingleDopplerParams, self).__init__()
        self.sonicFreq = 20000
        self.innerDirName = "/home/Dataset/final_data/ultra/12"
        self.algoType = AlgoType.SINGLE_DOPPLER
        self.volun = 1
        # self.trainDirList = ['/home/Dataset/all_ultra_data/1/20k_3cm_0d',
        #                      '/home/Dataset/all_ultra_data/1/20k_2cm',
        #                      '/home/Dataset/all_ultra_data/1/20k_4cm',
        #                      '/home/Dataset/all_ultra_data/1/20k_3cm_p10d_15',
        #                      '/home/Dataset/all_ultra_data/1/20k_3cm_n10d_15'
        #                      ]
        self.trainDirList = ['/home/Dataset/final_data/ultra/1',
                             '/home/Dataset/final_data/ultra/2',
                             '/home/Dataset/final_data/ultra/3',
                             '/home/Dataset/final_data/ultra/4',
                             '/home/Dataset/final_data/ultra/5',
                             '/home/Dataset/final_data/ultra/6',
                             '/home/Dataset/final_data/ultra/7',
                             '/home/Dataset/final_data/ultra/8',
                             '/home/Dataset/final_data/ultra/9',
                             '/home/Dataset/final_data/ultra/10',
                             ]
        self.testDirList = ['/home/Dataset/all_ultra_data/1/mix_3cm']
        self.isGeneratePath = True
        self.pathGenerateType = 'single_same'
        self.dataMode = DataModoType.MULTIPLE
        self.selectTrainNum = -1
        self.selectTestNum = -1


volunList = [16,17,18,19,20]
trainNumbers = [6,9,12,15,18,21,24]
modelAllType = ['resnet1']
modelAllFun = ['kl_attn']

dataParams = SingleDopplerParams()

yamlPath = 'config/conv.yaml'
file = open(yamlPath, 'r')
modelParams = yaml.load(file, Loader=yaml.FullLoader)
file.close()

checkPath = 'ModelWeights/volun_1_10_0.2833.pt'


for curVolun in volunList:
    volunPath = os.path.join('/home/Dataset/final_data/ultra', str(curVolun))
    dataParams.innerDirName = volunPath
    for curTrainNumber in trainNumbers:
        dataParams.selectTrainNum = curTrainNumber
        datasetGenerator = DatasetGenerator(dataParams)
        datasetGenerator.emsembleProcess()
        for modelType in modelAllType:
            for modelFunc in modelAllFun:
                modelParams['modelType'] = modelType
                modelParams['modelFunc'] = modelFunc
                modelParams['curVolun'] = curVolun
                modelParams['trainNumber'] = curTrainNumber
                modelParams['dataMode'] = dataParams.dataMode
                lipModel = LipReading(modelParams)
                lipModel.LoadCheck(checkPath)
                lipModel.Train()
    with open('all_result.txt', 'a+') as f:
        f.write('\n\n')